from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import time
from typing import Dict, Any
from core.aigc_detector import AIGCDetector

# 创建 FastAPI 应用实例
app = FastAPI(
    title="AIGC 图像伪造检测 API",
    description="计算 AIGC 检测模型的性能指标，包括 AP 和处理速度",
    version="1.0.0"
)


# 定义请求模型
class DetectionRequest(BaseModel):
    data_path: str
    # 可以添加其他可选参数，如模型配置等
    model: Dict[str, Any] = {"name":  "CLIP:ViT-L/14"}


# 定义响应模型
class DetectionResponse(BaseModel):
    ap: str  # 平均精度
    processing_speed: str  # 处理速度 (样本/秒)
    time_cost: float  # 耗时 (秒)
    model: Dict[str, Any]  # 模型配置


# 定义处理函数
def process_val(data_path: str, model: Dict[str, Any] = None):
    detector = AIGCDetector('ours')
    t0 = time.time()
    ap, len_data = detector.val(data_path)
    time_cost = time.time() - t0
    processing_speed = len_data / time_cost
    print(f'time_cost: {time_cost}s, processing_speed: {processing_speed} samples/s')
    return {
        "ap": str(round(ap, 4)),
        "processing_speed": str(round(processing_speed, 4)),
        "time_cost": time_cost,
        "model": model or {}
    }


# 定义 API 端点
@app.post("/detect", response_model=DetectionResponse, tags=["检测"])
async def detect_aigc(request: DetectionRequest):
    """
    计算 AIGC 检测模型的性能指标

    参数:
    - data_path: 数据集路径
    - model: 可选的模型配置参数

    返回:
    - ap: 平均精度
    - processing_speed: 处理速度 (样本/秒)
    - time_cost: 处理耗时 (秒)
    - model: 使用的模型配置
    """
    try:
        # 验证数据路径是否有效
        # 实际应用中可以添加更复杂的路径验证逻辑
        if not request.data_path:
            raise HTTPException(status_code=400, detail="数据路径不能为空")

        # 调用处理函数
        result = process_val(request.data_path, request.model)

        return result
    except Exception as e:
        # 异常处理
        raise HTTPException(status_code=500, detail=f"处理失败: {str(e)}")


# 启动应用
if __name__ == "__main__":
    import uvicorn
    uvicorn.run('detect:app', host="0.0.0.0", port=8080, reload=True)